// encodetest.cpp
// Copyright (C) 2008 Sam Bateman
// samuel.bateman@gmail.com
//
// This file is part of DVD On Fire.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#include <cstdio>
#include <cstring>
#include <ctime>
#include <stdexcept>
#include <string>
using namespace std;

// the libav* headers aren't c++ compatible
extern "C" {
#include <ffmpeg/avcodec.h>
#include <ffmpeg/avformat.h>
#include <ffmpeg/opt.h>
}

// struct AVData {{{
struct AVData
{
    AVData()
    : context(NULL), audio(NULL), video(NULL),
      audio_stream(NULL), video_stream(NULL), input(true) {}

    AVFormatContext *context;
    AVCodecContext  *audio;
    AVCodecContext  *video;
    AVStream        *audio_stream;
    AVStream        *video_stream;
    bool             input;
};
// }}}

// deleteAVData(data) {{{
void deleteAVData(AVData *data)
{
    if (data->input){
        avcodec_close(data->audio);
        avcodec_close(data->video);
        av_close_input_file(data->context);
    }
    else {
        av_free(data->context);
    }
}
// }}}
// openInputFile(file) {{{
AVData* openInputFile(const string& file)
{
    // open the input file
    AVData *data = new AVData;
    data->input = true;
    if (av_open_input_file(&data->context, file.c_str(), NULL, 0, NULL) != 0){
        throw runtime_error("Couldn't open file " + file);
    }
    if (av_find_stream_info(data->context) < 0){
        throw runtime_error("Couldn't get stream info in " + file);
    }

    // find the input video and audio streams
    for (unsigned int i = 0; i < data->context->nb_streams; ++i){
        AVStream *s = data->context->streams[i];
        AVCodecContext *c = s->codec;
        if (c->codec_type == CODEC_TYPE_VIDEO){
            data->video_stream = s;
        }
        else if (c->codec_type == CODEC_TYPE_AUDIO){
            data->audio_stream = s;
        }
    }
    if (data->video_stream == NULL || data->audio_stream == NULL){
        throw runtime_error("Missing video or audio stream");
    }
    data->video = data->video_stream->codec;
    data->audio = data->audio_stream->codec;

    // open the input video and audio codecs
    AVCodec *codec = avcodec_find_decoder(data->video->codec_id);
    if (codec == NULL || avcodec_open(data->video, codec) < 0){
        throw runtime_error("Couldn't open video codec");
    }
    codec = avcodec_find_decoder(data->audio->codec_id);
    if (codec == NULL || avcodec_open(data->audio, codec) < 0){
        throw runtime_error("Couldn't open audio codec");
    }

    // we're done
    dump_format(data->context, 0, file.c_str(), false);
    return data;
}
// }}}
// openOutputFile(file) {{{
AVData* openOutputFile(const string& file)
{
    // create the output file
    AVData *data = new AVData;
    data->input = false;
    data->context = av_alloc_format_context();
    data->context->oformat = guess_format("dvd", NULL, NULL);
    strcpy(data->context->filename, file.c_str());
    if (url_fopen(&data->context->pb, file.c_str(), URL_WRONLY) < 0){
        throw runtime_error("Couldn't open file " + file);
    }

    // create the video stream
    data->video_stream = av_new_stream(data->context, 0);
    data->video = data->video_stream->codec;
    data->video->codec_id = CODEC_ID_MPEG2VIDEO;
    data->video->codec_type = CODEC_TYPE_VIDEO;
    data->video->width = 720;
    data->video->height = 480;
    data->video->sample_aspect_ratio.num = 4;
    data->video->sample_aspect_ratio.den = 3;
    data->video->time_base.num = 1001;
    data->video->time_base.den = 30000;
    av_set_string(data->video, "b",       "6000000");
    av_set_string(data->video, "maxrate", "9000000");
    av_set_string(data->video, "minrate", "0"      );
    av_set_string(data->video, "bufsize", "1835008");

    // create the audio stream
    data->audio_stream = av_new_stream(data->context, 0);
    data->audio = data->audio_stream->codec;
    data->audio->codec_id = CODEC_ID_AC3;
    data->audio->codec_type = CODEC_TYPE_AUDIO;
    data->audio->sample_rate = 48000;
    data->audio->channels = 2;
    data->audio->time_base.num = 1;
    data->audio->time_base.den = data->audio->sample_rate;
    av_set_string(data->audio, "ab", "448000");

    // we're done
    dump_format(data->context, 1, file.c_str(), true);
    return data;
}
// }}}
// main(argc, argv) {{{
int main(int argc, char **argv)
{
    // check for command-line arguments
    if (argc < 2){
        printf("Usage: %s <mpg file>\n", argv[0]);
        return 1;
    }

    // start timing
    unsigned int start_time = time(NULL);

    // open the input and output files
    av_register_all();
    AVData *input = openInputFile(argv[1]);
    AVData *output = openOutputFile("etout.mpg");

    // read and write packets
    AVFrame *frame = avcodec_alloc_frame();
    AVPacket packet;
    int video_frames = 0;
    int audio_frames = 0;
    int video_packets = 0;
    int audio_packets = 0;
    while (av_read_frame(input->context, &packet) == 0){
        if (packet.stream_index == input->video_stream->index){
            int finished;
            avcodec_decode_video(input->video, frame, &finished,
                                 packet.data, packet.size);
            if (finished) ++video_frames;
            ++video_packets;
        }
        else if (packet.stream_index == input->audio_stream->index){
            ++audio_frames;
            ++audio_packets;
        }
        av_free_packet(&packet);
        printf("\r%d of 126378", video_frames);
        fflush(stdout);
    }
    printf("\n%d/%d video, %d/%d audio\n",
            video_frames, video_packets,
            audio_frames, audio_packets);

    // end timing
    unsigned int end_time = time(NULL);
    int time_taken = end_time - start_time;
    printf("Time taken: %dh %dm %ds\n",
           time_taken / 3600, time_taken / 60 % 60, time_taken % 60);

    // clean up
    av_free(frame);
    deleteAVData(input);
    deleteAVData(output);
}
// }}}
